// Copyright (c) 2013 The Chromium OS Authors. All rights reserved.
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

#include "memento_common.h"
#include "general_installer.h"

#include <sys/sendfile.h>
#include <sys/stat.h>
#include <sys/types.h>
#include <sys/wait.h>

#include <fcntl.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>

#include <iostream>
#include <sstream>

using std::cerr;
using std::endl;
using std::string;
using std::stringstream;

namespace chromeos_memento_updater {

GeneralInstaller::GeneralInstaller(const string& running_dir,
                                   const string& board) {
  running_dir_ = running_dir;
  board_ = board;
  combined_kernel_ = false;

  // Currently, we only allow partition options in image mode.
  allow_partition_options_ = false;
}

GeneralInstaller::~GeneralInstaller() {}

// TODO(chunyen) : convert the system call into C++ code
bool GeneralInstaller::IsDeviceMounted(const string& device_path) {
  FILE* fp;
  stringstream command;
  const int kMaxPathLen = 512;
  char buffer[kMaxPathLen];

  command << "grep \"^" << device_path <<
             " \" /proc/mounts | cut -d ' ' -f 1 | uniq";
  fp = popen(command.str().c_str(), "r");
  if (fp == NULL) {
    error_exit("Cannot check if device is mounted");
  }
  if (fgets(buffer, sizeof(buffer), fp) == NULL) {
    fclose(fp);
    return false;
  }
  fclose(fp);
  return (device_path == buffer);
}
int GeneralInstaller::PrepareImage(const int channel_id,
                                   FILE* out_pipe) {
  return error("This method should not be called, need override");
}

int GeneralInstaller::InstallImage(const string& device_path,
                                   const string& kernel_device_path,
                                   FILE* in_pipe) {
  stringstream command;

  if (in_pipe == NULL) {
    return error("Null input pipe");
  }

  // Unmount device when necessary
  if (IsDeviceMounted(device_path)) {
    if (umount(device_path.c_str()) != 0) {
      return error("Cannot umount " + device_path);
    }
  }

  command.str(string());
  if (kernel_device_path.empty()) {
    // TODO(chunyen) : Using dd instead of cat, or do a complete rewrite?
    command << "cat > " << device_path;
  } else {
    command << running_dir_ << "/split_write " <<
               kernel_device_path << " " << device_path;
  }
  FILE* out_pipe = popen(command.str().c_str(), "w");
  if (out_pipe == NULL) {
    return error("Fail to open device to write");
  }

  // Write by pipe or split write
  const int kBufSize = 1024 * 1024 * 4;  // 4MiB
  int read_byte, write_byte;
  char buf[kBufSize];
  while ((read_byte = fread(buf, 1, kBufSize, in_pipe)) > 0) {
    write_byte = 0;
    do {
      write_byte += fwrite(&buf[write_byte], 1,
                           read_byte - write_byte, out_pipe);
      if (ferror(out_pipe)) {
        ZeroOutDevice(device_path);
        return error("Error writing pipe");
      }
    } while (write_byte < read_byte);
  }
  pclose(out_pipe);
  if (ferror(in_pipe)) {
    ZeroOutDevice(device_path);
    return error("Error reading pipe");
  }
  // Flush linux cache.
  sync();

  return 0;
}

void GeneralInstaller::ResetPartitionTable(const string& device_path) {
  string pmbr = "/root/.pmbr_code";
  // Check if pmbr exists
  if (access(pmbr.c_str(), R_OK) != 0) {
    error_exit("Missing " + pmbr +", please rebuild image");
  }

  if (system((". /usr/sbin/write_gpt.sh &&"
              " write_base_table " + device_path + " " + pmbr).c_str()) != 0) {
    error_exit("Cannot write partition table.");
  }
  cerr << "Reloading partition table changes..."  << endl;
  sync();
  if (system(("partprobe " + device_path).c_str()) != 0) {
    error_exit("Cannot reload partition table.");
  }
}

void GeneralInstaller::InitStatefulPartition(const string& device_path) {
  if (system(("mkfs.ext4 " + device_path).c_str()) != 0) {
    error_exit("Cannot make stateful partition.");
  }
}

void GeneralInstaller::InstallDevice(const int channel_id,
                                     const string& device_path,
                                     const string& kernel_device_path) {
  int fd[2];
  pid_t pid;
  if (pipe(fd) != 0) {
    error_exit("Failed to create pipe");
  }
  pid = fork();
  if (pid < 0) {
    // Failed, exit.
    error_exit("Fork failed");
  } else if (pid == 0) {
    // Child process, close fd[1] (for write), run InstallImage and _Exit().
    close(fd[1]);
    FILE* fp = fdopen(fd[0], "rb");
    if (fp == NULL) {
      perror("Fail creating file pointer for read");
      _Exit(-1);
    }
    int return_value = InstallImage(device_path, kernel_device_path, fp);
    fclose(fp);
    _Exit(return_value);
  } else {
    // Parent process, close fd[0] (for read), run PrepareImage and wait().
    close(fd[0]);
    FILE* fp = fdopen(fd[1], "wb");
    if (fp == NULL) {
      error_exit("Fail creating file pointer for write");
    }
    int return_value = PrepareImage(channel_id, fp);
    fclose(fp);
    if (return_value != 0) {
      exit(return_value);
    }
    if (waitpid(pid, &return_value, 0) != pid) {
      error_exit("Wait pid failed");
    }
    if (return_value != 0) {
      exit(return_value);
    }
  }
}

void GeneralInstaller::ZeroOutDevice(const string& device_path) {
  system_call("dd if=/dev/zero of=" + device_path + " bs=4096 count=1",
              "Cannnot zero out device");
}

string GeneralInstaller::PrepareFirmware(const string& rootfs_path) {
  // Extract firmware updater from rootfs
  MountHandler mount_handler(rootfs_path, "ext2", MS_RDONLY);
  TempfileHandler temp_updater_handler(false);
  const string tmp_updater = temp_updater_handler.GetFilePath();
  const string updater_path = "/usr/sbin/chromeos-firmwareupdate";
  const string mount_point = mount_handler.GetMountPoint();

  // Here we use sendfile() to efficiently copy file without creating
  // extra buffer.
  int source = open((mount_point + updater_path).c_str(), O_RDONLY);
  int dest = open(tmp_updater.c_str(),
                  O_WRONLY | O_TRUNC | O_CREAT, S_IRWXU);
  struct stat stat_source;
  fstat(source, &stat_source);

  if (source == -1 || dest == -1) {
    if (source != -1) {
      close(source);
    }
    if (dest != -1) {
      close(dest);
    }
    error_exit("Fail to copy firmware updater");
  }
  if (sendfile(dest, source, 0, stat_source.st_size) == -1) {
    close(source);
    close(dest);
    error_exit("Fail to copy firmware updater");
  }
  close(source);
  close(dest);
  return tmp_updater;
}

void GeneralInstaller::UpdateFirmware(const string& updater_path,
                                      const string& mode) {
  stringstream command;

  command << "sh " << updater_path << " --force --mode=" << mode;
  system_call(command.str(), "Fail to run firmware updater");
}

void GeneralInstaller::ActivateImage(const string& device_path) {
  MountHandler handler(device_path, "ext2", MS_RDONLY);
  string mount_point = handler.GetMountPoint();
  system_call(mount_point + "/postinst " + device_path,
              "Cannot run postinst");
}

}  // namespace chromeos_memento_updater
